package com.example.raft;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;
import java.util.concurrent.locks.ReentrantLock;

/**
 * RaftNode
 *
 * 本类实现了简化版的 Raft 核心逻辑，包括：
 * 1. 领导者选举 (Leader Election)
 * 2. 日志复制 (Log Replication)
 * 3. 心跳、超时定时
 *
 * handleClientRequest(...) 模拟当本节点是 Leader 时处理客户端写请求；否则拒绝。
 */
public class RaftNode {

    /**
     * 节点的三种状态：FOLLOWER, CANDIDATE, LEADER
     */
    public enum NodeState {
        FOLLOWER,
        CANDIDATE,
        LEADER
    }

    /**
     * LogKey 用于标识一条日志的 (term, index)
     */
    public static class LogKey {
        private final long term;  // 日志所在任期
        private final long index; // 日志索引

        public LogKey(long term, long index) {
            this.term = term;
            this.index = index;
        }

        public long getTerm() {
            return term;
        }

        public long getIndex() {
            return index;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            LogKey that = (LogKey) o;
            return term == that.term && index == that.index;
        }

        @Override
        public int hashCode() {
            return Objects.hash(term, index);
        }

        @Override
        public String toString() {
            return String.format("(term=%d, index=%d)", term, index);
        }
    }

    /**
     * 日志条目：term, index, command
     */
    public static class LogEntry {
        private final long term;     // 日志任期
        private final long index;    // 日志索引
        private final String command;// 具体指令(或数据)

        public LogEntry(long term, long index, String command) {
            this.term = term;
            this.index = index;
            this.command = command;
        }

        public long getTerm() {
            return term;
        }

        public long getIndex() {
            return index;
        }

        public String getCommand() {
            return command;
        }

        @Override
        public String toString() {
            return String.format("LogEntry(term=%d,index=%d,cmd=%s)", term, index, command);
        }
    }

    // -------------- 关键字段 --------------

    protected final String nodeId;                     // 本节点 ID
    protected final AtomicReference<NodeState> state;  // 当前节点状态
    protected final AtomicLong currentTerm;            // 当前任期
    protected final AtomicReference<String> votedFor;  // 本任期内投票给哪位

    /**
     * logMap 存放日志：key=(term, index)，value=LogEntry
     */
    protected final ConcurrentHashMap<LogKey, LogEntry> logMap;

    /**
     * commitIndex：已经被多数节点复制的最高日志索引
     */
    protected final AtomicLong commitIndex;

    /**
     * lastApplied：已应用到状态机的最高日志索引
     */
    protected final AtomicLong lastApplied;

    /**
     * 分配日志索引时使用的自增计数器
     */
    protected final AtomicLong lastLogIndexAllocated;

    /**
     * scheduler：定时线程池，用来执行选举超时和心跳等
     */
    protected final ScheduledExecutorService scheduler;

    /**
     * 两个定时任务：选举超时 / 心跳定时
     */
    protected ScheduledFuture<?> electionTimeoutFuture;
    protected ScheduledFuture<?> heartbeatFuture;

    /**
     * lock：保护对状态、term 等关键变量的并发更新
     */
    protected final ReentrantLock lock;

    /**
     * peerIds：集群中其它节点ID
     */
    protected final List<String> peerIds;

    // 常量：RPC重试次数，选举尝试次数，节点个数
    private static final int MAX_RPC_RETRIES = 3;
    private static final int MAX_ELECTION_ROUNDS = 2;
    private static final int NODE_COUNT = 10;

    // candidateElectionRetryCount：在同一个term内，Candidate多次请求投票的次数
    private int candidateElectionRetryCount = 0;

    // -------------- 构造 --------------

    public RaftNode(String nodeId, List<String> peerIds) {
        this.nodeId = nodeId;
        this.state = new AtomicReference<>(NodeState.FOLLOWER);
        this.currentTerm = new AtomicLong(0);
        this.votedFor = new AtomicReference<>(null);

        this.logMap = new ConcurrentHashMap<>();
        this.commitIndex = new AtomicLong(0);
        this.lastApplied = new AtomicLong(0);
        this.lastLogIndexAllocated = new AtomicLong(0);

        // 有几个节点，就创建几条线程的pool用于选举超时、心跳等调度
        this.scheduler = Executors.newScheduledThreadPool(NODE_COUNT);
        this.lock = new ReentrantLock();
        this.peerIds = (peerIds != null) ? peerIds : Collections.emptyList();

        // 启动时立即设置选举超时
        resetElectionTimeout();
    }

    // -------------- 选举超时 --------------

    /**
     * resetElectionTimeout：重置/重新设置选举超时
     * 若已有定时任务，则取消之
     */
    protected void resetElectionTimeout() {
        if (electionTimeoutFuture != null) {
            electionTimeoutFuture.cancel(true);
        }
        // 随机 150~300 ms
        int timeoutMillis = ThreadLocalRandom.current().nextInt(150, 301);
        // 到时间后调用 startElection
        electionTimeoutFuture = scheduler.schedule(this::startElection, timeoutMillis, TimeUnit.MILLISECONDS);
    }

    /**
     * startElection：当Follower长时间未收到心跳，或Candidate超时未当选，会调用
     */
    protected void startElection() {
        lock.lock();
        try {
            // 若已是Leader，无需选举
            if (state.get() == NodeState.LEADER) {
                return;
            }
            // 转成Candidate
            state.set(NodeState.CANDIDATE);
            // 任期 + 1
            long newTerm = currentTerm.incrementAndGet();
            // 给自己投票
            votedFor.set(nodeId);

            // 选举轮次清0
            candidateElectionRetryCount = 0;
            System.out.println(nodeId + " becomes CANDIDATE, term=" + newTerm);
        } finally {
            lock.unlock();
        }

        // 并行请求投票
        requestVotesFromPeers();
    }

    /**
     * requestVotesFromPeers：向其他节点并发发送 RequestVoteRPC
     * 若一轮不成功且没超 MAX_ELECTION_ROUNDS，则再尝试
     */
    protected void requestVotesFromPeers() {
        final long termSnapshot = currentTerm.get();
        final AtomicInteger voteCount = new AtomicInteger(1); // 自己1票
        final CountDownLatch latch = new CountDownLatch(peerIds.size());

        long myLastIndex = getLastLogIndex();
        long myLastTerm = getLastLogTerm();

        for (String peerId : peerIds) {
            // 用线程池异步发送
            scheduler.execute(() -> {
                boolean granted = false;
                // 多次RPC重试
                for (int i = 0; i < MAX_RPC_RETRIES; i++) {
                    if (tryRequestVoteRPC(peerId, termSnapshot, nodeId, myLastIndex, myLastTerm)) {
                        granted = true;
                        break;
                    }
                }
                if (granted) {
                    voteCount.incrementAndGet();
                }
                latch.countDown();
            });
        }

        // 等待300ms或全部完成
        try {
            latch.await(300, TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        // 汇总投票结果
        lock.lock();
        try {
            // 如果期间term变了，作废
            if (termSnapshot < currentTerm.get()) {
                return;
            }
            // 是否得到多数票
            if (voteCount.get() > (peerIds.size() / 2)) {
                becomeLeader();
            } else {
                // 未获多数票
                candidateElectionRetryCount++;
                if (candidateElectionRetryCount < MAX_ELECTION_ROUNDS) {
                    System.out.println("Candidate " + nodeId + " no majority, attempt="
                            + candidateElectionRetryCount + "/" + MAX_ELECTION_ROUNDS
                            + ", term=" + currentTerm.get());
                    requestVotesFromPeers();
                } else {
                    System.out.println("Candidate " + nodeId + " gave up election, wait next timeout.");
                    resetElectionTimeout();
                }
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * tryRequestVoteRPC：对 requestVoteRPC 的一次封装
     */
    private boolean tryRequestVoteRPC(String peerId, long term, String candidateId,
                                      long lastIndex, long lastTerm) {
        return requestVoteRPC(peerId, term, candidateId, lastIndex, lastTerm);
    }

    /**
     * becomeLeader：若当选Leader，取消选举定时，启动心跳
     */
    protected void becomeLeader() {
        state.set(NodeState.LEADER);
        if (electionTimeoutFuture != null) {
            electionTimeoutFuture.cancel(true);
        }
        System.out.println(nodeId + " becomes LEADER, term=" + currentTerm.get());
        startHeartbeat();
    }

    // -------------- 心跳 --------------

    /**
     * startHeartbeat：定期(50ms)给Follower发心跳
     */
    protected void startHeartbeat() {
        if (heartbeatFuture != null && !heartbeatFuture.isCancelled()) {
            heartbeatFuture.cancel(true);
        }
        heartbeatFuture = scheduler.scheduleAtFixedRate(() -> {
            if (state.get() == NodeState.LEADER) {
                sendHeartbeat();
            }
        }, 0, 50, TimeUnit.MILLISECONDS);
    }

    /**
     * sendHeartbeat：实际发 AppendEntries（空日志）到各peer
     */
    protected void sendHeartbeat() {
        long leaderTerm = currentTerm.get();
        long prevLogIndex = getLastLogIndex();
        long prevLogTerm = getLastLogTerm();

        for (String peerId : peerIds) {
            scheduler.execute(() -> {
                boolean success = false;
                for (int i = 0; i < MAX_RPC_RETRIES; i++) {
                    if (appendEntriesRPC(peerId, leaderTerm, nodeId,
                            prevLogIndex, prevLogTerm, null, null, commitIndex.get())) {
                        success = true;
                        break;
                    }
                }
                if (!success) {
                    System.out.println("[Heartbeat] " + nodeId + " => peer " + peerId + " failed after retries");
                }
            });
        }
    }

    /**
     * requestVoteRPC：模拟发送投票请求到对方节点
     */
    protected boolean requestVoteRPC(String peerId, long term, String candidateId,
                                     long lastIndex, long lastTerm) {
        RaftNode peer = SimulationEnv.getNode(peerId);
        if (peer == null) return false;
        return peer.onRequestVote(term, candidateId, lastIndex, lastTerm);
    }

    /**
     * onRequestVote：处理对方发来的RequestVote
     */
    public boolean onRequestVote(long term, String candidateId, long candidateLastIndex, long candidateLastTerm) {
        lock.lock();
        try {
            // 对方term大于我，则更新本地term
            if (term > currentTerm.get()) {
                currentTerm.set(term);
                state.set(NodeState.FOLLOWER);
                votedFor.set(null);
                if (heartbeatFuture != null) heartbeatFuture.cancel(true);
                resetElectionTimeout();
            }
            // 若对方term依旧比我小，拒绝
            if (term < currentTerm.get()) {
                return false;
            }
            // 如果本节点已投票给别人且不是 candidateId，拒绝
            String myVote = votedFor.get();
            if (myVote != null && !myVote.equals(candidateId)) {
                return false;
            }
            // 检查对方日志是否足够新
            long myLastIndex = getLastLogIndex();
            long myLastTerm = getLastLogTerm();
            boolean upToDate = (candidateLastTerm > myLastTerm)
                    || (candidateLastTerm == myLastTerm && candidateLastIndex >= myLastIndex);
            if (!upToDate) {
                return false;
            }
            // 投票给对方
            votedFor.set(candidateId);
            resetElectionTimeout();
            return true;
        } finally {
            lock.unlock();
        }
    }

    // -------------- 日志复制 --------------

    /**
     * handleClientRequest：模拟客户端请求写入一条日志(仅Leader可成功)
     */
    public boolean handleClientRequest(String command) {
        if (state.get() != NodeState.LEADER) {
            System.out.println("Node " + nodeId + " is not leader, reject command: " + command);
            return false;
        }
        lock.lock();
        try {
            // 分配新的日志index
            long newIndex = lastLogIndexAllocated.incrementAndGet();
            long curTerm = currentTerm.get();
            LogKey key = new LogKey(curTerm, newIndex);
            LogEntry entry = new LogEntry(curTerm, newIndex, command);

            // 先写到本地
            logMap.put(key, entry);
            System.out.println("Leader " + nodeId + " appended log: " + entry);

            // 并行发 AppendEntries 给其他节点
            final AtomicInteger successCount = new AtomicInteger(1);
            final CountDownLatch latch = new CountDownLatch(peerIds.size());

            for (String peerId : peerIds) {
                scheduler.execute(() -> {
                    boolean success = false;
                    for (int i = 0; i < MAX_RPC_RETRIES; i++) {
                        if (appendEntriesRPC(peerId, curTerm, nodeId,
                                getPrevLogIndex(newIndex), getPrevLogTerm(newIndex),
                                key, entry, commitIndex.get())) {
                            success = true;
                            break;
                        }
                    }
                    if (success) {
                        successCount.incrementAndGet();
                    }
                    latch.countDown();
                });
            }

            // 等待300ms
            try {
                latch.await(300, TimeUnit.MILLISECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }

            // 若过半成功，则commit
            if (successCount.get() > (peerIds.size() / 2)) {
                long oldCommit = commitIndex.get();
                commitIndex.set(newIndex);
                applyEntries(oldCommit + 1, newIndex);
                return true;
            } else {
                // 否则回退
                logMap.remove(key);
                return false;
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * onAppendEntries：Follower 处理 Leader 的日志复制请求(或心跳)
     */
    public boolean onAppendEntries(long leaderTerm, String leaderId,
                                   long prevLogIndex, long prevLogTerm,
                                   LogKey newLogKey, LogEntry newLogEntry,
                                   long leaderCommit) {
        lock.lock();
        try {
            // 若对方term更大，更新自己为Follower
            if (leaderTerm > currentTerm.get()) {
                currentTerm.set(leaderTerm);
                state.set(NodeState.FOLLOWER);
                votedFor.set(null);
                if (heartbeatFuture != null) heartbeatFuture.cancel(true);
                resetElectionTimeout();
            }
            if (leaderTerm < currentTerm.get()) {
                return false;
            }
            // 确保是Follower
            if (state.get() != NodeState.FOLLOWER) {
                state.set(NodeState.FOLLOWER);
            }
            resetElectionTimeout();

            // 如果是空日志(心跳)
            if (newLogKey == null || newLogEntry == null) {
                long myCommit = commitIndex.get();
                if (leaderCommit > myCommit) {
                    long newCommit = Math.min(leaderCommit, getLastLogIndex());
                    commitIndex.set(newCommit);
                    applyEntries(myCommit + 1, newCommit);
                }
                return true;
            }

            // 检查prevLogIndex
            if (prevLogIndex > 0) {
                LogKey prevKey = new LogKey(prevLogTerm, prevLogIndex);
                if (!logMap.containsKey(prevKey)) {
                    removeConflictLogs(prevLogIndex);
                    return false;
                }
            }
            // 判断是否覆盖
            LogEntry existing = logMap.get(newLogKey);
            if (existing != null) {
                if (existing.getTerm() != newLogEntry.getTerm()) {
                    // 冲突，删除后面的
                    removeConflictLogs(newLogKey.getIndex());
                    logMap.put(newLogKey, newLogEntry);
                }
            } else {
                logMap.put(newLogKey, newLogEntry);
                long allocated = lastLogIndexAllocated.get();
                if (newLogKey.getIndex() > allocated) {
                    lastLogIndexAllocated.set(newLogKey.getIndex());
                }
            }

            // 更新commitIndex
            long myCommit = commitIndex.get();
            if (leaderCommit > myCommit) {
                long newCommit = Math.min(leaderCommit, getLastLogIndex());
                commitIndex.set(newCommit);
                applyEntries(myCommit + 1, newCommit);
            }
            return true;
        } finally {
            lock.unlock();
        }
    }

    /**
     * appendEntriesRPC：模拟网络调用 peer.onAppendEntries
     * 并加10%随机失败
     */
    protected boolean appendEntriesRPC(String peerId, long leaderTerm, String leaderId,
                                       long prevLogIndex, long prevLogTerm,
                                       LogKey newLogKey, LogEntry newLogEntry,
                                       long leaderCommit) {
        RaftNode peer = SimulationEnv.getNode(peerId);
        if (peer == null) return false;

        // 10% 概率随机失败
        if (ThreadLocalRandom.current().nextInt(100) < 10) {
            return false;
        }
        return peer.onAppendEntries(leaderTerm, leaderId, prevLogIndex, prevLogTerm,
                newLogKey, newLogEntry, leaderCommit);
    }

    /**
     * removeConflictLogs：从 conflictIndex 开始的日志都删除
     */
    protected void removeConflictLogs(long conflictIndex) {
        for (LogKey k : logMap.keySet()) {
            if (k.getIndex() >= conflictIndex) {
                logMap.remove(k);
            }
        }
        long maxIndex = getLastLogIndex();
        lastLogIndexAllocated.set(maxIndex);
    }

    /**
     * applyEntries：将 [startIndex..endIndex] 范围的日志应用到本地状态机(这里仅打印)
     */
    protected void applyEntries(long startIndex, long endIndex) {
        for (long i = startIndex; i <= endIndex; i++) {
            long idx = i;
            Optional<LogKey> maybeKey = logMap.keySet().stream()
                    .filter(k -> k.getIndex() == idx).findFirst();
            if (maybeKey.isPresent()) {
                LogEntry e = logMap.get(maybeKey.get());
                if (e != null) {
                    persistLogEntryToDB(e);
                    lastApplied.incrementAndGet();
                }
            }
        }
    }

    /**
     * persistLogEntryToDB：模拟写数据库，这里仅打印
     */
    protected void persistLogEntryToDB(LogEntry entry) {
        System.out.println("Node " + nodeId + " => [DB Write] " + entry);
    }

    // -------------- 工具方法 --------------

    protected long getLastLogIndex() {
        return logMap.keySet().stream().mapToLong(LogKey::getIndex).max().orElse(0);
    }

    protected long getLastLogTerm() {
        long idx = getLastLogIndex();
        if (idx == 0) return 0;
        Optional<LogKey> maybeKey = logMap.keySet().stream()
                .filter(k -> k.getIndex() == idx).findFirst();
        if (maybeKey.isPresent()) {
            return maybeKey.get().getTerm();
        }
        return 0;
    }

    protected long getPrevLogIndex(long newIndex) {
        return (newIndex - 1) <= 0 ? 0 : (newIndex - 1);
    }

    protected long getPrevLogTerm(long newIndex) {
        long prevIndex = newIndex - 1;
        if (prevIndex <= 0) return 0;
        Optional<LogKey> maybeKey = logMap.keySet().stream()
                .filter(k -> k.getIndex() == prevIndex).findFirst();
        if (maybeKey.isPresent()) {
            return maybeKey.get().getTerm();
        }
        return 0;
    }

    /**
     * stopNode：停止所有定时任务和线程
     */
    public void stopNode() {
        if (electionTimeoutFuture != null) {
            electionTimeoutFuture.cancel(true);
        }
        if (heartbeatFuture != null) {
            heartbeatFuture.cancel(true);
        }
        scheduler.shutdownNow();
        System.out.println("Node " + nodeId + " stopped.");
    }

    // -------------- Getter --------------
    public String getNodeId() {
        return nodeId;
    }

    public NodeState getState() {
        return state.get();
    }

    public long getCurrentTerm() {
        return currentTerm.get();
    }

    public long getCommitIndex() {
        return commitIndex.get();
    }

    public long getLastApplied() {
        return lastApplied.get();
    }
}